Hugging Faceのモデル学習で、ロス関数をカスタマイズする方法

nokomoro3

2022.11.15

この記事は公開されてから1年以上経過しています。情報が古い可能性がありますので、ご注意ください。

こんちには。

データアナリティクス事業本部機械学習チームの中村です。

Hugging Faceのライブラリの使い方紹介記事第５弾です。

今回は、学習時にロス関数をカスタマイズする方法をご紹介します。

ロス関数をカスタマイズする目的

ロス関数をカスタマイズする目的には、例えば以下のようなものが挙げられます。

不均衡なデータに対するアプローチの１つ
- 例えば、実際のタスクに取り組む歳は不均衡なデータなども取り扱うシーンが発生します。
- 不均衡なデータに対するアプローチはいくつかありますが、ロス関数に手を入れるというのも一つの手法です。
デフォルトと異なるロス関数を使用したい
- 単純にデフォルトとは異なるロス関数に異なるものを使いたい場合もあります。
- Hugging FaceでTrainerを使用する場合は、デフォルトはnn.CrossEntropyLossとなります。

今回はこのような場合に必要となるロス関数のカスタマイズについて説明します。

例として、不均衡なデータに対応するために、クラスごとの重みづけを実施してロスを計算する実装を行います。

実行環境

今回はGoogle Colaboratory環境で実行しました。

ハードウェアなどの情報は以下の通りです。

GPU: Tesla T4 (GPUメモリ16GB搭載)
CUDA: 11.2
メモリ: 13GB

主なライブラリのバージョンは以下となります。

transformers: 4.24.0
datasets: 2.6.1

インストール

transformersとdatasetsをインストールします。

!pip install transformers datasets

また事前学習モデルの依存モジュールをインストールします。

!pip install fugashi
!pip install ipadic
!pip install sentencepiece

ベースとするコード

今回のベースとするコードは以下のとおりです。

import random
from datasets import load_dataset, DatasetDict
from transformers import AutoTokenizer
from transformers import AutoModelForSequenceClassification
from transformers import TrainingArguments
from transformers import Trainer
from sklearn.metrics import accuracy_score, f1_score
import torch

# データセットのロード
dataset = load_dataset("tyqiangz/multilingual-sentiments", "japanese")

# 実験のためデータセットを縮小したい場合はコチラを有効化
random.seed(42)
dataset = DatasetDict({
    "train"     : dataset['train']\
        .select(random.sample(range(dataset['train']     .num_rows), k=1000)),
    "validation": dataset['validation']\
        .select(random.sample(range(dataset['validation'].num_rows), k=1000)),
    "test"      : dataset['test']\
        .select(random.sample(range(dataset['test']      .num_rows), k=1000)),
})

# トークナイザのロード
model_ckpt = "cl-tohoku/bert-base-japanese-whole-word-masking"
tokenizer = AutoTokenizer.from_pretrained(model_ckpt)

# トークナイズ処理
def tokenize(batch):
    return tokenizer(batch["text"], padding=True, truncation=True)
dataset_encoded = dataset.map(tokenize, batched=True, batch_size=None)

# 事前学習モデルのロード
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
num_labels = 3
model = (AutoModelForSequenceClassification
    .from_pretrained(model_ckpt, num_labels=num_labels)
    .to(device))

# メトリクスの定義
def compute_metrics(pred):
    labels = pred.label_ids
    preds = pred.predictions.argmax(-1)
    f1 = f1_score(labels, preds, average="weighted")
    acc = accuracy_score(labels, preds)
    return {"accuracy": acc, "f1": f1}

# 学習パラメータの設定
batch_size = 16
logging_steps = len(dataset_encoded["train"]) // batch_size
model_name = "sample-text-classification-bert"

training_args = TrainingArguments(
    output_dir=model_name,
    num_train_epochs=2,
    learning_rate=2e-5,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    weight_decay=0.01,
    evaluation_strategy="epoch",
    disable_tqdm=False,
    logging_steps=logging_steps,
    push_to_hub=False,
    log_level="error",
)

# Trainerの定義
trainer = Trainer(
    model=model,
    args=training_args,
    compute_metrics=compute_metrics,
    train_dataset=dataset_encoded["train"],
    eval_dataset=dataset_encoded["validation"],
    tokenizer=tokenizer
)

# トレーニング実行
trainer.train()

この内容についての解説は以下の記事を参照ください。

またコード中に記載の通り、実験のみを簡易に試されたい方は以下を有効にして試されてください。

逆にフルサイズのデータで試したい方は以下をコメントアウトして使用ください。

# 実験のためデータセットを縮小したい場合はコチラを有効化
random.seed(42)
dataset = DatasetDict({
    "train"     : dataset['train']\
        .select(random.sample(range(dataset['train']     .num_rows), k=1000)),
    "validation": dataset['validation']\
        .select(random.sample(range(dataset['validation'].num_rows), k=1000)),
    "test"      : dataset['test']\
        .select(random.sample(range(dataset['test']      .num_rows), k=1000)),
})

今回は有効化した前提で記事を書きます。

CustomTrainerの作成

公式の例

ロス関数をカスタムするには、Trainerクラスを独自に定義します。

以下にTrainerのサブクラスを使ってカスタマイズするAPIについて記載があります。

このページにちょうど、ロス関数をカスタムする例の記載があります。

class CustomTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        labels = inputs.get("labels")
        # forward pass
        outputs = model(**inputs)
        logits = outputs.get("logits")
        # compute custom loss (suppose one has 3 labels with different weights)
        loss_fct = nn.CrossEntropyLoss(weight=torch.tensor([1.0, 2.0, 3.0]))
        loss = loss_fct(logits.view(-1, self.model.config.num_labels), labels.view(-1))
        return (loss, outputs) if return_outputs else loss

少しこの記述について解説をします。

関数の引数で与えられるinputsは辞書型で、このケースの場合以下のキーを持つようです。

dict_keys(['input_ids', 'token_type_ids', 'attention_mask', 'labels'])

モデルの出力であるoutputsも、このケースの場合以下のキーを持つようです。

odict_keys(['loss', 'logits'])

logitsは、ミニバッチの各クラスのlogitですので、(batch_size, クラス数)の形状となります。

こちらを参考にして、データセットのラベルの比率で重みづけをしてみます。

ラベルの比率を使った重みづけ

DataFrameに変換し、ラベルの比率を見てみます。

dataset.set_format(type="pandas")
label_ratio = dataset["train"][:].value_counts("label", sort=False, normalize=True)
dataset.reset_format()
print(label_ratio)

label
0    0.338
1    0.319
2    0.343
dtype: float64

偏りはそこまでありませんが、今回はこの逆数を重みづけに使ってみます。

(逆数では激しすぎる場合などは対数を使うものありかと思います)

逆数にするついでにPyTorchで使用できるようなtensor型、device設定に変換します。

class_weight = torch.tensor(1/label_ratio).clone().to(device, torch.float32)
class_weight

tensor([2.9586, 3.1348, 2.9155], device='cuda:0')

こちらの重みづけをCustomTrainer内のnn.CrossEntropyLossに与えます。

from torch import nn

class CustomTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        labels = inputs.get("labels")
        # forward pass
        outputs = model(**inputs)
        logits = outputs.get("logits")
        # compute custom loss (suppose one has 3 labels with different weights)
        loss_fct = nn.CrossEntropyLoss(weight=class_weight)
        loss = loss_fct(logits.view(-1, self.model.config.num_labels), labels.view(-1))
        return (loss, outputs) if return_outputs else loss

以下のコードで再度学習します。Trainerの代わりにCustomTrainerを使っています。

# モデル再定義
model = (AutoModelForSequenceClassification
    .from_pretrained(model_ckpt, num_labels=num_labels)
    .to(device))

# Trainerの定義
trainer = CustomTrainer(
    model=model,
    args=training_args,
    compute_metrics=compute_metrics,
    train_dataset=dataset_encoded["train"],
    eval_dataset=dataset_encoded["validation"],
    tokenizer=tokenizer
)

# トレーニング実行
trainer.train()

これでロス関数をカスタマイズした学習ができました。